package org.micode.libs.logtrace;

import org.micode.common.logtrace.LogTracer;

import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class TraceThreadPoolExecutor extends ThreadPoolExecutor {

    protected final Map<Runnable, String> traceIdMap = new ConcurrentHashMap<>();

    public TraceThreadPoolExecutor(int corePoolSize,
                                   int maxPoolSize,
                                   int keepAliveSeconds,
                                   TimeUnit seconds,
                                   BlockingQueue<Runnable> queue,
                                   ThreadFactory threadFactory,
                                   RejectedExecutionHandler rejectedExecutionHandler) {
        super(corePoolSize, maxPoolSize, keepAliveSeconds, seconds, queue, threadFactory, rejectedExecutionHandler);
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
        LogTracer.setTraceId(traceIdMap.get(r));
        super.beforeExecute(t, r);
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
        traceIdMap.remove(r);
        try {
            super.afterExecute(r, t);
        } finally {
            LogTracer.clearTraceId();
        }
    }

    @Override
    protected void terminated() {
        super.terminated();
    }

    @Override
    public void execute(Runnable command) {
        traceIdMap.put(command, LogTracer.getTraceId());
        super.execute(command);
    }

}
